import os
import pandas as pd
import random
from sklearn.model_selection import train_test_split
import shutil

def create_test_set(input_folder, output_folder, test_size=0.2, random_state=None):
    """
    从包含多个CSV文件的文件夹中随机选取一定比例的文件作为测试集，
    并将其复制到新的测试集文件夹中，同时将剩余文件作为训练集复制到相应的文件夹中。

    参数:
    input_folder: 包含原始CSV文件的文件夹路径
    output_folder: 存储测试集和训练集CSV文件的新文件夹路径
    test_size: 测试集占总文件数的比例，默认为0.2（即20%）
    random_state: 随机种子，保证结果可复现，默认为None
    """
    # 设置随机种子以确保结果可复现
    if random_state is not None:
        random.seed(random_state)

    # 确保输出文件夹存在
    for folder in ['test', 'train']:
        path = os.path.join(output_folder, folder)
        if not os.path.exists(path):
            os.makedirs(path)

    # 获取所有CSV文件
    csv_files = [f for f in os.listdir(input_folder) if f.endswith('.csv')]

    # 计算需要选择的文件数量
    num_files = len(csv_files)
    num_test_files = int(num_files * test_size)

    # 随机选择一定比例的文件作为测试集
    test_files = set(random.sample(csv_files, num_test_files))
    train_files = set(csv_files) - test_files

    # 将选中的文件复制到测试集文件夹，其余文件复制到训练集文件夹
    for file_name in csv_files:
        source_path = os.path.join(input_folder, file_name)
        if file_name in test_files:
            destination_path = os.path.join(output_folder, 'test', file_name)
        else:
            destination_path = os.path.join(output_folder, 'train', file_name)
        shutil.copy(source_path, destination_path)
        print(f"Copied {file_name} to {destination_path}")

# 使用函数
input_folder = r"D:\temp\xinda\ModelRecognition\总数据集"  # 替换为你的CSV文件夹路径
output_folder = r"D:\temp\xinda\ModelRecognition\随机选取"  # 替换为你想要保存测试集和训练集的文件夹路径
create_test_set(input_folder, output_folder, test_size=0.2)  # 设置测试集大小为10%